TurboQuant encoding for Vectors by lwwmanning · Pull Request #7167 · vortex-data/vortex

lwwmanning · 2026-03-25T21:35:19Z

Summary

Lossy quantization for vector data (e.g., embeddings) based on TurboQuant

TODOs & Follow-ups

Scalar function dispatch wiring

L2Norm readthrough: l2_norm_direct() in compute/l2_norm.rs returns the stored norms child directly (exact, O(1) per vector). Needs to be wired into the vortex-tensor L2Norm scalar function dispatch so that l2_norm(Extension(TurboQuant(...))) short-circuits without decompression. Options: register an encoding-specific kernel in initialize(), or add a TurboQuant check in the tensor crate's L2Norm executor.
CosineSimilarity in quantized domain: cosine_similarity_quantized() in compute/cosine_similarity.rs computes approximate cosine similarity via centroid lookup in the rotated domain — no full decompression needed. Needs wiring into vortex-tensor cosine_similarity dispatch when both arguments come from the same TurboQuant column (same rotation and codebook). Accuracy is bounded by the quantization distortion (~O(1/4^b)).

Compression pipeline

Pluggable compressor integration: The BtrBlocks builder integration (with_turboquant) is deferred pending @connortsui20's pluggable compressor work. Currently TurboQuant is invoked directly via compress_turboquant() in the canonical compressor.

Feature gaps

Nullable input support: The encoder currently rejects nullable FixedSizeListArray inputs. To support nullable vectors, we'd need to strip validity before encoding and reapply it after (e.g., via MaskedArray wrapper), or propagate it through a validity child.

codspeed-hq · 2026-03-25T21:39:27Z

Merging this PR will not alter performance

✅ 1106 untouched benchmarks
🆕 20 new benchmarks
⏩ 1522 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	`BASE`	`HEAD`	Efficiency
🆕	Simulation	`turboquant_compress_dim1024_2bit_qjl`	N/A	137.4 ms	N/A
🆕	Simulation	`turboquant_compress_dim128_4bit_mse`	N/A	13.2 ms	N/A
🆕	Simulation	`turboquant_compress_dim768_4bit_qjl`	N/A	161.9 ms	N/A
🆕	Simulation	`turboquant_compress_dim128_4bit_qjl`	N/A	20.7 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_2bit_mse`	N/A	37.1 ms	N/A
🆕	Simulation	`turboquant_compress_dim768_4bit_mse`	N/A	103.4 ms	N/A
🆕	Simulation	`turboquant_compress_dim1024_2bit_mse`	N/A	67.3 ms	N/A
🆕	Simulation	`turboquant_compress_dim1024_4bit_mse`	N/A	104.3 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_4bit_mse`	N/A	37.1 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_4bit_qjl`	N/A	83.1 ms	N/A
🆕	Simulation	`turboquant_decompress_dim128_4bit_qjl`	N/A	10.7 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_2bit_qjl`	N/A	83.1 ms	N/A
🆕	Simulation	`turboquant_compress_dim1024_8bit_qjl`	N/A	482.8 ms	N/A
🆕	Simulation	`turboquant_decompress_dim128_4bit_mse`	N/A	4.9 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_8bit_qjl`	N/A	83 ms	N/A
🆕	Simulation	`turboquant_decompress_dim1024_8bit_mse`	N/A	37.1 ms	N/A
🆕	Simulation	`turboquant_compress_dim1024_8bit_mse`	N/A	731.2 ms	N/A
🆕	Simulation	`turboquant_decompress_dim768_4bit_qjl`	N/A	78.6 ms	N/A
🆕	Simulation	`turboquant_compress_dim1024_4bit_qjl`	N/A	163 ms	N/A
🆕	Simulation	`turboquant_decompress_dim768_4bit_mse`	N/A	36.6 ms	N/A

_{Comparing claude/admiring-lichterman (df5c8f9) with develop (4e4ea56)}

1522 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

vortex-tensor/src/encodings/turboquant/vtable.rs

AdamGS · 2026-03-27T12:36:38Z

encodings/turboquant/src/centroids.rs

+/// Maximum iterations for Max-Lloyd algorithm.
+const MAX_ITERATIONS: usize = 200;
+
+type CentroidCache = Mutex<HashMap<(u32, u8), Vec<f32>>>;


should really be a dashmap

encodings/turboquant/src/centroids.rs

vortex-tensor/src/encodings/turboquant/centroids.rs

encodings/turboquant/src/centroids.rs

encodings/turboquant/src/compress.rs

vortex-tensor/src/encodings/turboquant/mod.rs

vortex-btrblocks/src/builder.rs

connortsui20 · 2026-03-27T13:22:03Z

Do we want to put this in vortex-tensor?

Implement the TurboQuant algorithm (arXiv:2504.19874) as a new lossy encoding for high-dimensional vector data. This supports both the MSE-optimal and inner-product-optimal (Prod) variants at 1-4 bits per coordinate. Key components: - Max-Lloyd centroid computation on Beta(d/2,d/2) distribution - Deterministic random rotation via nalgebra QR decomposition - FastLanes BitPackedArray for index storage - QJL residual correction for unbiased inner product estimation (Prod) The encoding operates on FixedSizeList arrays of floats, which is the storage format for Vector and FixedShapeTensor extension types. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

…ntegration Add a CompressorPlugin wrapper that intercepts Vector and FixedShapeTensor extension columns, applies TurboQuant encoding, and recursively compresses the resulting children (norms, codes) via the inner compressor. Expose this via WriteStrategyBuilder::with_vector_quantization(config), which composes with existing encoding modes (default, compact, cuda). TODO: restructure into BtrBlocks canonical_compressor directly (like DateTimeParts) rather than the wrapper CompressorPlugin approach. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Move TurboQuant compression logic from a standalone CompressorPlugin wrapper into the BtrBlocks canonical compressor, following the same pattern as DateTimeParts. This gives TurboQuant access to the full BtrBlocks recursive compression pipeline for its children (norms, codes, etc.). Changes: - Add `turboquant_config: Option<TurboQuantConfig>` to BtrBlocksCompressor - Add `with_turboquant(config)` to BtrBlocksCompressorBuilder - Add tensor extension detection + compress_turboquant() in the Canonical::Extension arm of canonical_compressor - Update WriteStrategyBuilder::with_vector_quantization to configure BtrBlocks directly instead of wrapping - Remove TurboQuantCompressor wrapper and vortex-layout dep from vortex-turboquant Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Add TurboQuant benchmarks to the single_encoding_throughput suite, covering compress and decompress for dim=128 and dim=768 at 2-bit and 4-bit widths. Uses 1000 random N(0,1) vectors per benchmark. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

…nsform Replace the O(d²) dense matrix rotation (previously nalgebra, then faer) with a Structured Random Hadamard Transform (SRHT) that runs in O(d log d). The SRHT applies D₃·H·D₂·H·D₁ where H is the Walsh-Hadamard transform and Dₖ are random diagonal ±1 sign matrices. This eliminates both the nalgebra and faer dependencies — the SRHT is fully self-contained with no external linear algebra library needed. Benchmark results (1000 vectors, mean throughput): | Benchmark | Before (nalgebra) | After (SRHT) | |----------------------------|---------:|----------:| | compress dim128 2-bit | 222 MB/s | 242 MB/s | | compress dim768 2-bit | 32 MB/s | 181 MB/s | | decompress dim128 2-bit | 87 MB/s | 614 MB/s | | decompress dim768 2-bit | 6 MB/s | 458 MB/s | For non-power-of-2 dimensions (e.g., 768), input is zero-padded to the next power of 2 (1024) and all padded coordinates are quantized. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

…tests Replace the loose "normalized MSE < 1.0" check with rigorous tests: - mse_within_theoretical_bound: Verifies per-vector normalized MSE is within 10x the paper's Theorem 1 bound (sqrt(3)*pi/2 / 4^b). Tests across dim={128,256} x bits={1,2,3,4}. - prod_inner_product_bias: Verifies the Prod variant produces approximately unbiased inner products by computing <query, x_hat> vs <query, x> over 500 random pairs and checking mean relative error < 0.3. - mse_decreases_with_bits: Verifies MSE monotonically decreases with increasing bit-width for both Mse and Prod variants. Total: 49 tests (up from 39). Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

- Hoist per-row allocations (residual, projected) out of encode_prod loop - Use BufferMut<u8> directly for sign_buf instead of Vec + copy - Remove unused num-traits dependency - Remove dead unreachable!() branch (bit_width >= 2 validated at entry) - Fix orphaned doc comment blank line - Generate public-api.lock files for new/modified crates Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Address code review findings: - Tighten SRHT roundtrip test tolerance from 1e-3 to 1e-5 (verified exact to ~4e-7 relative error across dim 32-1024). Consolidate into parameterized rstest covering power-of-2 and non-power-of-2 dims. - Rename `pd` -> `padded_dim` throughout compress.rs and decompress.rs for clarity. - Add early dimension validation (>= 2) in turboquant_encode with clear error message. - Add edge case tests: single-row roundtrip (Mse + Prod), empty array Prod variant, dimension-below-2 rejection. - Tighten norm preservation test to 1e-5 relative tolerance. Total: 59 tests (up from 49). Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

…ror bounds Add comprehensive crate documentation including: - Theoretical MSE bounds per bit-width from the paper's Theorem 1 - Compression ratio table for common dimensions (256-1536), accounting for power-of-2 padding overhead on non-power-of-2 dims (768, 1536) - Working doctest demonstrating encode usage and size verification Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Extend bit_width range from 1-4 to 1-8. At 8 bits (256 centroids), codes are stored as raw u8 instead of bit-packed since BitPackedArray doesn't support width >= 8. This gives ~4x compression from f32 with near-lossless quality (MSE bound 4.15e-05). Changes: - Update all validation sites (compress, array, centroids) to accept 1-8 bits (MSE) and 2-8 bits (Prod) - Skip bitpack_encode for 8-bit codes, store PrimitiveArray<u8> directly - Extend crate docs with full 1-8 bit bound/ratio tables - Add 6-bit and 8-bit test cases for roundtrip, MSE bounds, Prod bias, and monotonic MSE decrease. High bit-width tests verify MSE < 4-bit MSE and MSE < 1% (since the theoretical bound becomes unrealistically tight at 5+ bits due to SRHT finite-dimension effects) - Regenerate public-api.lock Total: 69 unit tests + 1 doctest. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Allow Prod variant bit_width up to 9, where the MSE component uses 8-bit codes (raw u8) plus 1-bit QJL correction. The 8-bit MSE codes can be fed directly into int8 GEMM kernels on tensor cores without unpacking. - Update Prod validation to 2-9, MSE remains 1-8 - Restructure top-level validation into per-variant match - Add 9-bit roundtrip, inner product bias, and monotonicity tests - Document tensor core use case in crate docs Total: 71 unit tests + 1 doctest. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Expand TurboQuant throughput benchmarks to cover common embedding dimensions: - dim=128 (2-bit, 4-bit) — small embeddings - dim=768 (2-bit) — BERT / sentence-transformers - dim=1024 (2-bit, 4-bit) — larger embedding models - dim=1536 (2-bit, 4-bit) — OpenAI ada-002, exercises non-power-of-2 padding overhead All benchmarks use i.i.d. N(0,1) vectors with fixed seed — a conservative worst-case for TurboQuant since real neural embeddings have structure that the SRHT exploits for better quantization. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Add methods to persist and restore SRHT rotation signs as BoolArray, eliminating the need to regenerate from seed during decompression: - `export_inverse_signs_bool_array()`: Exports 3 × padded_dim sign bits as a single BoolArray in inverse-application order [D₃|D₂|D₁] so decompression iterates sequentially. - `from_bool_array(signs, dim)`: Reconstructs RotationMatrix from stored signs without needing the seed. - `apply_inverse_srht_from_bits(buf, signs_bytes, padded_dim, norm_factor)`: Hot-path free function that applies inverse SRHT directly from raw sign bytes, avoiding intermediate Vec<f32> reconstruction. Convention: bit=1 means +1, bit=0 means -1 (negate). Tests verify: - Export→import roundtrip produces identical rotation (3 dims) - Hot-path function matches struct-based inverse_rotate exactly Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Add two new cascading array types that replace the monolithic TurboQuantArray: TurboQuantMSEArray (4 children): - codes (BitPackedArray or PrimitiveArray<u8>) - norms (PrimitiveArray<f32>) - centroids (PrimitiveArray<f32>, stored codebook) - rotation_signs (BoolArray, 3 * padded_dim bits, inverse order) TurboQuantQJLArray (4 children): - mse_inner (TurboQuantMSEArray at bit_width - 1) - qjl_signs (BoolArray, num_rows * padded_dim) - residual_norms (PrimitiveArray<f32>) - rotation_signs (BoolArray, QJL rotation, inverse order) Both store all precomputed data (centroids, rotation signs) as children to eliminate recomputation during decompression. Validity is pushed down to the codes child via ValidityVTableFromChild at each level. Includes decompression implementations for both new types that use stored centroids/signs and the hot-path apply_inverse_srht_from_bits. The old TurboQuantArray and its decode paths are retained for now. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Add `turboquant_encode_mse()` and `turboquant_encode_qjl()` that produce the new cascaded array types: - turboquant_encode_mse: produces TurboQuantMSEArray with stored centroids (PrimitiveArray<f32>) and rotation signs (BoolArray) - turboquant_encode_qjl: produces TurboQuantQJLArray wrapping an inner TurboQuantMSEArray at bit_width-1, with QJL signs (BoolArray) and QJL rotation signs (BoolArray) Tests verify: - Roundtrip encode/decode for both new types at various dims/bit_widths - New MSE path matches legacy path exactly (bit-for-bit) - Edge cases: empty arrays and single-row arrays for both types Total: 90 unit tests + 1 doctest. Signed-off-by: Will Manning <will@spiraldb.com> Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Will Manning <will@willmanning.io>

Signed-off-by: Will Manning <will@willmanning.io>

This reverts commit 0c5e8e73af9afc001e20405c91d11d59a8129796. Signed-off-by: Will Manning <will@willmanning.io>

Signed-off-by: Will Manning <will@willmanning.io>

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com> Signed-off-by: Will Manning <will@willmanning.io>

…o vortex-tensor Signed-off-by: Will Manning <will@willmanning.io>

…o vortex-tensor (pt 2) Signed-off-by: Will Manning <will@willmanning.io>

Signed-off-by: Will Manning <will@willmanning.io>

…chterman

Signed-off-by: Will Manning <will@willmanning.io>

This reverts commit a928727.

Signed-off-by: Will Manning <will@willmanning.io>

This reverts commit 00ee4fe.

This reverts commit 54b158c.

Signed-off-by: Will Manning <will@willmanning.io>

…chterman

lwwmanning added the changelog/feature A new feature label Mar 26, 2026